package com.duyulong.demo.controller;

import com.duyulong.demo.service.ChatService;
import com.duyulong.demo.repository.UserRepository;
import com.duyulong.demo.entity.User;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.security.core.Authentication;
import org.springframework.stereotype.Controller;
import org.springframework.ui.Model;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestMethod;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.ResponseBody;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import javax.servlet.http.HttpSession;
import java.util.Optional;

@Controller
public class ChatController {
    private static final Logger logger = LoggerFactory.getLogger(ChatController.class);

    @Autowired
    private ChatService chatService;

    @Autowired
    private ObjectMapper objectMapper;

    @Autowired
    private UserRepository userRepository;

    @RequestMapping(value = "/chat", method = RequestMethod.GET)
    public String chatPage(Model model, Authentication authentication, HttpSession session) {
        if (authentication != null) {
            String username = authentication.getName();
            model.addAttribute("username", username);
            
            // 首先尝试从session获取用户信息
            String nickname = (String) session.getAttribute("nickname");
            String email = (String) session.getAttribute("email");
            
            // 如果session中没有信息，从数据库获取
            if (nickname == null) {
                Optional<User> userOptional = userRepository.findByUsername(username);
                if (userOptional.isPresent()) {
                    User user = userOptional.get();
                    nickname = user.getNickname();
                    email = user.getEmail();
                    // 更新session
                    session.setAttribute("nickname", nickname);
                    session.setAttribute("email", email);
                }
            }
            
            // 设置模型属性
            if (nickname != null) {
                model.addAttribute("nickname", nickname);
            }
            if (email != null) {
                model.addAttribute("email", email);
            }
        }
        
        return "chat";
    }

    @RequestMapping(value = "/api/chat/stream", method = RequestMethod.POST)
    public ResponseEntity<StreamingResponseBody> streamChat(@RequestParam String message, Authentication authentication) {
        if (authentication == null) {
            return ResponseEntity.status(401)
                .contentType(MediaType.TEXT_EVENT_STREAM)
                .body(outputStream -> {
                    ObjectNode eventNode = objectMapper.createObjectNode();
                    eventNode.put("code", 401);
                    eventNode.put("msg", "请先登录后再使用");
                    eventNode.put("time", System.currentTimeMillis());
                    writeEvent(outputStream, eventNode.toString());
                    writeEvent(outputStream, "[DONE]");
                });
        }

        logger.info("开始处理流式请求，消息内容长度: {}", message.length());
        
        StreamingResponseBody responseBody = outputStream -> {
            try {
                logger.info("开始调用chatService.streamChat处理消息");
                chatService.streamChat(message, new ChatService.StreamCallback() {
                    @Override
                    public void onMessage(String content) {
                        try {
                            if (content == null || content.isEmpty()) {
                                logger.info("收到空内容，跳过处理");
                                return;
                            }
                            
                            logger.info("收到消息内容，长度: {}", content.length());
                            ObjectNode eventNode = objectMapper.createObjectNode();
                            eventNode.put("code", 200);
                            eventNode.put("time", System.currentTimeMillis());
                            
                            ObjectNode deltaNode = objectMapper.createObjectNode();
                            deltaNode.put("content", content);
                            deltaNode.put("role", "assistant");
                            
                            ObjectNode choiceNode = objectMapper.createObjectNode();
                            choiceNode.put("index", 0);
                            choiceNode.set("delta", deltaNode);
                            choiceNode.putNull("finish_reason");
                            
                            eventNode.putArray("choices").add(choiceNode);
                            eventNode.put("created", System.currentTimeMillis() / 1000);
                            eventNode.put("id", "chatglm_" + System.currentTimeMillis());
                            eventNode.put("object", "chat.completion.chunk");
                            
                            writeEvent(outputStream, eventNode.toString());
                            outputStream.flush();
                            logger.info("消息内容已发送并刷新");
                        } catch (IOException e) {
                            logger.error("发送消息内容时发生IO异常", e);
                            Thread.currentThread().interrupt();
                        }
                    }

                    @Override
                    public void onError(String error) {
                        try {
                            logger.error("处理消息时发生错误: {}", error);
                            ObjectNode eventNode = objectMapper.createObjectNode();
                            eventNode.put("code", 500);
                            eventNode.put("msg", error);
                            eventNode.put("time", System.currentTimeMillis());
                            writeEvent(outputStream, eventNode.toString());
                            writeEvent(outputStream, "[DONE]");
                            outputStream.flush();
                            logger.info("错误消息已发送并刷新");
                        } catch (IOException e) {
                            logger.error("发送错误消息时发生IO异常", e);
                            Thread.currentThread().interrupt();
                        }
                    }

                    @Override
                    public void onComplete() {
                        try {
                            logger.info("开始发送完成消息");
                            ObjectNode eventNode = objectMapper.createObjectNode();
                            eventNode.put("code", 200);
                            eventNode.put("time", System.currentTimeMillis());
                            
                            ObjectNode choiceNode = objectMapper.createObjectNode();
                            choiceNode.put("index", 0);
                            choiceNode.put("finish_reason", "stop");
                            
                            ObjectNode deltaNode = objectMapper.createObjectNode();
                            deltaNode.put("role", "assistant");
                            choiceNode.set("delta", deltaNode);
                            
                            eventNode.putArray("choices").add(choiceNode);
                            eventNode.put("created", System.currentTimeMillis() / 1000);
                            eventNode.put("id", "chatglm_" + System.currentTimeMillis());
                            eventNode.put("object", "chat.completion.chunk");
                            
                            writeEvent(outputStream, eventNode.toString());
                            writeEvent(outputStream, "[DONE]");
                            outputStream.flush();
                            logger.info("完成消息已发送并刷新");
                        } catch (IOException e) {
                            logger.error("发送完成消息时发生IO异常", e);
                            Thread.currentThread().interrupt();
                        }
                    }
                });
            } catch (Exception e) {
                try {
                    logger.error("处理消息时发生异常", e);
                    ObjectNode eventNode = objectMapper.createObjectNode();
                    eventNode.put("code", 500);
                    eventNode.put("msg", "服务器内部错误：" + e.getMessage());
                    eventNode.put("time", System.currentTimeMillis());
                    writeEvent(outputStream, eventNode.toString());
                    writeEvent(outputStream, "[DONE]");
                    outputStream.flush();
                    logger.info("异常消息已发送并刷新");
                } catch (IOException ignored) {
                    logger.error("发送异常消息时发生IO异常", ignored);
                    Thread.currentThread().interrupt();
                }
            }
        };

        logger.info("准备返回StreamingResponseBody");
        return ResponseEntity.ok()
                .contentType(MediaType.TEXT_EVENT_STREAM)
                .header("Cache-Control", "no-cache")
                .header("Connection", "keep-alive")
                .header("X-Accel-Buffering", "no")
                .body(responseBody);
    }

    private void writeEvent(OutputStream outputStream, String data) throws IOException {
        if (Thread.currentThread().isInterrupted()) {
            logger.warn("检测到线程已被中断，停止写入");
            throw new IOException("Stream closed");
        }
        
        synchronized (outputStream) {
            try {
                outputStream.write(("data: " + data + "\n\n").getBytes(StandardCharsets.UTF_8));
                logger.trace("成功写入数据，长度: {}", data.length());
            } catch (IOException e) {
                if (e.getMessage() != null && 
                    (e.getMessage().contains("Broken pipe") || 
                     e.getMessage().contains("Connection reset"))) {
                    logger.warn("检测到连接已断开: {}", e.getMessage());
                    Thread.currentThread().interrupt();
                    return;
                }
                logger.error("写入数据时发生IO异常", e);
                throw e;
            }
        }
    }

    @RequestMapping(value = "/api/chat", method = RequestMethod.POST)
    @ResponseBody
    public String chat(@RequestParam String message) {
        return chatService.chat(message);
    }
} 